"""
图像校正模块

用于检测和校正图像的倾斜、旋转等问题
提供自动角度检测和图像矫正功能
"""

import cv2
import numpy as np
from skimage.transform import rotate
from skimage.feature import canny
from skimage.transform import hough_line, hough_line_peaks


class ImageCorrection:
    """
    图像校正处理类
    
    功能：
    - 自动检测图像倾斜角度
    - 校正图像旋转和倾斜
    - 优化图像对比度和亮度
    """
    
    def __init__(self):
        """初始化图像校正器"""
        self.max_angle = 45  # 最大校正角度
    
    def detect_skew_angle(self, image: np.ndarray) -> float:
        """
        检测图像倾斜角度
        
        参数:
            image: 输入图像数组
            
        返回:
            float: 检测到的倾斜角度（度）
            
        使用场景:
            corrector = ImageCorrection()
            angle = corrector.detect_skew_angle(image_array)
        """
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image
        
        # 边缘检测
        edges = canny(gray, low_threshold=50, high_threshold=150)
        
        # 霍夫变换检测直线
        tested_angles = np.linspace(-np.pi/4, np.pi/4, 180, endpoint=False)
        h, theta, d = hough_line(edges, theta=tested_angles)
        
        # 找到最强的直线
        hough_peaks = hough_line_peaks(h, theta, d, num_peaks=20)
        
        if len(hough_peaks[1]) == 0:
            return 0.0
        
        # 计算主要角度
        angles = []
        for angle in hough_peaks[1]:
            angles.append(np.degrees(angle))
        
        # 找到最常见的角度
        angle_hist, bins = np.histogram(angles, bins=36, range=(-45, 45))
        peak_angle_idx = np.argmax(angle_hist)
        peak_angle = (bins[peak_angle_idx] + bins[peak_angle_idx + 1]) / 2
        
        return peak_angle
    
    def correct_skew(self, image: np.ndarray, angle: float = None) -> np.ndarray:
        """
        校正图像倾斜
        
        参数:
            image: 输入图像数组
            angle: 指定校正角度，如果为None则自动检测
            
        返回:
            numpy.ndarray: 校正后的图像
            
        使用场景:
            corrected = corrector.correct_skew(image_array)
            # 或指定角度
            corrected = corrector.correct_skew(image_array, angle=5.2)
        """
        if angle is None:
            angle = self.detect_skew_angle(image)
        
        # 限制角度范围
        angle = np.clip(angle, -self.max_angle, self.max_angle)
        
        if abs(angle) < 0.5:  # 角度太小，不需要校正
            return image
        
        # 旋转图像
        rotated = rotate(image, -angle, resize=True, preserve_range=True)
        
        # 确保数据类型正确
        return rotated.astype(np.uint8)
    
    def enhance_contrast(self, image: np.ndarray) -> np.ndarray:
        """
        增强图像对比度
        
        参数:
            image: 输入图像数组
            
        返回:
            numpy.ndarray: 对比度增强后的图像
            
        使用场景:
            enhanced = corrector.enhance_contrast(image_array)
        """
        # 转换为LAB色彩空间
        if len(image.shape) == 3:
            lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
            l_channel = lab[:, :, 0]
        else:
            l_channel = image
        
        # 应用CLAHE（限制对比度自适应直方图均衡）
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced_l = clahe.apply(l_channel)
        
        if len(image.shape) == 3:
            lab[:, :, 0] = enhanced_l
            enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        else:
            enhanced = enhanced_l
        
        return enhanced
    
    def auto_correct(self, image: np.ndarray) -> np.ndarray:
        """
        自动校正图像（包括倾斜校正和对比度增强）
        
        参数:
            image: 输入图像数组
            
        返回:
            numpy.ndarray: 自动校正后的图像
            
        使用场景:
            corrector = ImageCorrection()
            result = corrector.auto_correct(image_array)
        """
        # 首先校正倾斜
        corrected = self.correct_skew(image)
        
        # 然后增强对比度
        enhanced = self.enhance_contrast(corrected)
        
        return enhanced 